package com.yundingai.common;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriComponents;
import org.springframework.web.util.UriComponentsBuilder;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
public class DrawingWebSocketHandler extends TextWebSocketHandler {

    // 保存所有连接的会话，按房间分组
    private final Map<String, Set<WebSocketSession>> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String roomId = getRoomIdFromQuery(session);
        sessions.computeIfAbsent(roomId, k -> new CopyOnWriteArraySet<>()).add(session);
        System.out.println("用户加入房间：" + roomId);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String roomId = getRoomIdFromQuery(session);
        Set<WebSocketSession> roomSessions = sessions.get(roomId);
        if (roomSessions != null) {
            roomSessions.remove(session);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String roomId = getRoomIdFromQuery(session);
        // 广播给房间内所有用户
        for (WebSocketSession s : sessions.getOrDefault(roomId, Collections.emptySet())) {
            if (s.isOpen() && s != session) {
                s.sendMessage(message);
            }
        }
    }

    private String getRoomIdFromQuery(WebSocketSession session) {
        UriComponents uri = UriComponentsBuilder.fromUri(session.getUri()).build();
        return uri.getQueryParams().getFirst("roomId");
    }
}